import numpy as np

from .kalman_filter import KalmanFilter
from .track import Track
from .data_association import associate_detections_to_trackers


def x1y1x2y2_to_cxcyah(bboxes):
    new_boxes = np.zeros_like(bboxes)
    new_boxes[:, 0] = (bboxes[:, 0] + bboxes[:, 2]) / 2
    new_boxes[:, 1] = (bboxes[:, 1] + bboxes[:, 3]) / 2
    new_boxes[:, 2] = (bboxes[:, 2] - bboxes[:, 0]) / \
        (bboxes[:, 3] - bboxes[:, 1] + 1e-5)  # add 1e-5 in case divide by zero
    new_boxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
    return new_boxes


class SORT(object):
    def __init__(self, n_init=3, max_age=10, iou_threshold=0.3):
        self.kf = KalmanFilter()
        self.tracks = list()

        self.n_init = n_init
        self.max_age = max_age
        self.iou_threshold = iou_threshold
        self._next_id = 1

    def track(self, det_bboxes: np.ndarray) -> np.ndarray:
        """
        Parameters
        ----------
        det_bboxes : np.ndarray
            detected boxes, [x1, y1, x2, y2] format

        Returns
        -------
        np.ndarray
            return an numpy array, [track_id, x1, y1, x2, y2] format
        """
        predict_bboxes = self._predict()
        matches, unmatched_detections, unmatched_trackers = \
            associate_detections_to_trackers(
                det_bboxes, predict_bboxes, self.iou_threshold)

        det_bboxes = x1y1x2y2_to_cxcyah(det_bboxes)
        for match_pair in matches:
            det_index, track_index = match_pair
            self.tracks[track_index].update(self.kf, det_bboxes[det_index])
        for i in unmatched_detections:
            self._initiate_track(det_bboxes[i])
        for i in unmatched_trackers:
            self.tracks[i].mark_missed()
        self.tracks = [t for t in self.tracks if not t.is_deleted()]
        return np.array([[track.track_id] + track.to_tlbr().tolist()
                         for track in self.tracks])

    def _predict(self):
        trks = np.zeros((len(self.tracks), 4))
        to_del = list()
        for t, trk in enumerate(trks):
            self.tracks[t].predict(self.kf)
            trks[t] = self.tracks[t].to_tlbr()
            if np.any(np.isnan(trks[t])):
                to_del.append(t)
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
        for t in reversed(to_del):
            self.trackers.pop(t)
        return trks

    def _initiate_track(self, bboxes: np.ndarray):
        """
        Parameters
        ----------
        bboxes : np.ndarray
            detected boxes, [c_x, c_y, a, h] format
        """
        mean, covariance = self.kf.initiate(bboxes)
        self.tracks.append(
            Track(mean, covariance, self._next_id, self.n_init, self.max_age))
        self._next_id += 1
